from zhipuai import ZhipuAI
from typing import Dict, List, Optional
import json

class DialysisAIAssistant:
    def __init__(self, api_key: str):
        self.client = ZhipuAI(api_key=api_key)
        self.model = "glm-4-0520"
        
    def get_system_prompt(self) -> str:
        return """你是一位专业的血液透析顾问，具有以下专业能力：
        1. 解答血透相关专业问题
        2. 分析透析数据并给出建议
        3. 提供日常护理和饮食指导
        4. 解释医学术语
        5. 识别潜在风险并提供建议
        
        回答要求：
        - 使用通俗易懂的语言
        - 严格基于医学事实
        - 遇到紧急情况要建议就医
        - 重要信息要标记加粗"""

    def analyze_data(self, data: Dict) -> str:
        """分析透析数据"""
        messages = [
            {"role": "system", "content": self.get_system_prompt()},
            {"role": "user", "content": f"请分析这组透析数据并给出建议：{json.dumps(data, ensure_ascii=False)}"}
        ]
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages
        )
        return response.choices[0].message.content

    def chat_stream(self, query: str, context: Optional[Dict] = None) -> str:
        """流式对话"""
        messages = [
            {"role": "system", "content": self.get_system_prompt()},
            {"role": "user", "content": query}
        ]
        
        if context:
            messages.append({
                "role": "user",
                "content": f"参考患者当前数据：{json.dumps(context, ensure_ascii=False)}"
            })
        
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            stream=True
        )
        
        for chunk in response:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

def main():
    assistant = DialysisAIAssistant(api_key="your_api_key")
    
    # 测试数据分析
    test_data = {
        "血压": "110/70",
        "心率": 75,
        "超滤率": 12,
        "血流量": 250,
        "透析时间": 120,
        "体重": 65
    }
    
    print("=== 数据分析示例 ===")
    result = assistant.analyze_data(test_data)
    print(result)
    
    print("\n=== 对话示例 ===")
    test_questions = [
        "透析期间血压低怎么办？",
        "内瘘有什么护理方法？",
        "透析病人饮食上要注意什么？"
    ]
    
    for question in test_questions:
        print(f"\n问题: {question}")
        assistant.chat_stream(question)
        print("\n")

# if __name__ == "__main__":
#     main()